import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten, Dropout
from tensorflow.keras.datasets import mnist
from tensorflow.keras.utils import to_categorical
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import gradio as gr
import numpy as np
from PIL import Image

# 加载MNIST数据集并进行预处理
(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train = x_train.reshape(-1, 28 * 28).astype("float32") / 255
x_test = x_test.reshape(-1, 28 * 28).astype("float32") / 255
y_train = to_categorical(y_train, 10)
y_test = to_categorical(y_test, 10)

# 构建MLP模型
model = Sequential([
    Dense(512, input_dim=28*28, activation='relu'),
    Dropout(0.2),
    Dense(256, activation='relu'),
    Dropout(0.2),
    Dense(10, activation='softmax')
])

# 编译模型
model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])

# 训练模型
model.fit(x_train, y_train, epochs=10, batch_size=200, validation_split=0.1)

# 保存模型
model.save('best_mlp.h5')

# 加载保存的MLP模型
loaded_model = tf.keras.models.load_model('best_mlp.h5')

# 定义预处理函数
def preprocess(image):
    image = Image.fromarray(image).convert('L')  # 转换为灰度图像
    image_array = np.array(image).reshape((1, 28 * 28))  # 转换为一维数组并添加批次维度
    image_array = image_array.astype("float32") / 255  # 归一化
    return image_array

# 定义预测函数
def predict(image):
    preprocessed_image = preprocess(image)
    predicted_digit = loaded_model.predict(preprocessed_image)
    return np.argmax(predicted_digit)  # 返回最可能的数字

# 创建Gradio接口
iface = gr.Interface(fn=predict, inputs='sketchpad', outputs='text')

# 启动Gradio接口
iface.launch()